import os
import pandas as pd
from langchain_community.document_loaders import UnstructuredWordDocumentLoader, UnstructuredMarkdownLoader, PyPDFLoader


def save_result(dataset, file_path):
    """
    保存测试集生成结果
    """
    # 转为pandas格式
    dp = dataset.to_pandas()
    # 获取reference_contexts列表中最大长度
    max_len = dp['reference_contexts'].str.len().max()
    # 将reference_contexts列表拆分成多个列
    for i in range(max_len):
        dp[f'参考上下文{i+1}'] = [
            lst[i] if i < len(lst) else None 
            for lst in dp['reference_contexts']
        ]
    # 删除原始的reference_contexts、synthesizer_name列
    dp = dp.drop('reference_contexts', axis=1)
    dp = dp.drop('synthesizer_name', axis=1)
    # 重命名列
    dp.rename(columns={
        'user_input': '问题',
        'reference': '参考答案'
    }, inplace=True)
    
    # 使用通用方法保存
    write_to_excel(dp, file_path)


def save_eval_result(repr_dict: dict, file_name: str, file_path: str):
    """
    保存评估结果到excel
    """
    # 评估指标中英文映射
    metrics_map = {
        'llm_context_precision_without_reference': '上下文精度',
        'llm_context_precision_with_reference': '上下文精度(带参考)',
        'non_llm_context_precision_with_reference': '上下文精度(非LLM)',
        'context_recall': '上下文召回率',
        'non_llm_context_recall': '上下文召回率(非LLM)',
        'noise_sensitivity_relevant': '噪声敏感性',
        'answer_relevancy': '响应相关性',
        'faithfulness': '忠实度'
    }
    # 准备数据
    data = {'文件名': [file_name]}
    # 遍历所有可能的指标，只添加存在的指标
    for metric_name, metric_cn in metrics_map.items():
        if metric_name in repr_dict:
            data[metric_cn] = [repr_dict[metric_name]]
    # 转换为DataFrame
    df = pd.DataFrame(data)
    # 使用通用方法保存（追加模式）
    write_to_excel(df, file_path)


def create_dir_if_not_exists(file_path):
    """
    创建目录

    Args:
        file_path: 文件路径
    """
    dir_path = os.path.dirname(file_path)
    os.makedirs(dir_path, exist_ok=True)


def get_loader(file_path):
    """
    获取文件加载器

    Args:
        file_path: 文件路径
    Returns:
        loader: 文件加载器
    """
    ext = os.path.splitext(file_path)[1].lower()
    if ext == '.pdf':
        return PyPDFLoader(file_path)
    elif ext == '.md':
        return UnstructuredMarkdownLoader(file_path)
    elif ext in ['.docx', '.doc']:
        return UnstructuredWordDocumentLoader(file_path)
    else:
        raise ValueError(f"不支持的文件类型: {ext}")


def read_excel_file(file_path: str) -> pd.DataFrame:
    """
    读取excel文件

    Args:
        file_path: 文件路径
    Returns:
        pd.DataFrame: 数据框，如果读取失败则返回空DataFrame
    """
    try:
        return pd.read_excel(file_path)
    except Exception as e:
        print(f"读取文件 {file_path} 失败: {str(e)}")
        return pd.DataFrame()  # 返回空DataFrame而不是None


def write_to_excel(df: pd.DataFrame, file_path: str, index: bool = False, mode: str = 'w') -> None:
    """
    通用的Excel写入方法
    
    Args:
        df: 要写入的DataFrame
        file_path: 保存路径
        index: 是否保存索引，默认False
        mode: 写入模式，'w'为覆盖，'a'为追加，默认'w'
    """
    # 确保保存目录存在
    create_dir_if_not_exists(file_path)
    try:
        if mode == 'a' and os.path.exists(file_path):
            # 追加模式且文件存在
            existing_df = pd.read_excel(file_path)
            df = pd.concat([existing_df, df], ignore_index=True)
        
        # 保存到excel
        df.to_excel(file_path, index=index)
        print(f"数据已保存到: {file_path}")
    except Exception as e:
        print(f"保存文件失败: {str(e)}")


def get_filename_without_ext(file_path: str) -> str:
    """
    获取不带后缀的文件名
    
    Args:
        file_path: 文件路径或文件名
    Returns:
        str: 不带后缀的文件名
    """
    return os.path.splitext(os.path.basename(file_path))[0]


def get_files_to_process(specific_files: list[str] | None = None, file_path: str | None = None) -> list[str]:
    """
    获取需要处理的文件列表
    
    Args:
        specific_files: 指定的文件列表，如果为None则处理file_path目录下的所有文件
        file_path: 文档目录路径，当specific_files为None时使用
    Returns:
        list[str]: 需要处理的文件路径列表
    """
    files_to_process = []
    
    if specific_files:
        # 处理指定的文件列表
        for file in specific_files:
            if os.path.isfile(file):
                files_to_process.append(file)
            else:
                print(f"警告: 指定的文件 {file} 不存在，将被跳过")
    elif file_path:
        # 处理目录下的所有文件
        files_to_process = [
            os.path.join(file_path, f) 
            for f in os.listdir(file_path) 
            if os.path.isfile(os.path.join(file_path, f))
        ]
    
    if not files_to_process:
        print("没有找到需要处理的文件")
        
    return files_to_process

